package info.wangyuan.agent.config;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import info.wangyuan.agent.common.Constant.ModelConstant;
import info.wangyuan.agent.common.Enum.ModelType;
import info.wangyuan.agent.entity.po.ModelConfig;
import info.wangyuan.agent.mapper.ModelConfigMapper;
import org.springframework.stereotype.Component;

import java.time.Duration;


/**
 * 模型工厂
 *
 * @author Albert
 * @since 2025-07-09 22:15:26
 */
@Component
public class ModelFactory {

    private final ModelConfigMapper modelConfigMapper;

    public ModelFactory(ModelConfigMapper modelConfigMapper) {
        this.modelConfigMapper = modelConfigMapper;
    }

    /**
     * 创建对话模型
     */
    public ChatModel getCustomChatModel() {
        ModelConfig config = modelConfigMapper.selectOne(
                Wrappers.lambdaQuery(ModelConfig.class)
                        .eq(ModelConfig::getType, ModelType.LLM.getValue())
                        .eq(ModelConfig::getStatus, ModelConstant.IS_ACTIVE)
        );
        return new OpenAiChatModel.OpenAiChatModelBuilder()
                .strictJsonSchema(true)
                .logRequests(true)
                .logResponses(true)
                .returnThinking(false)
                .timeout(Duration.ofMinutes(3))
                .temperature(config.getTemperature())
                .apiKey(config.getApiKey().trim())
                .modelName(config.getName().trim())
                .baseUrl(config.getBaseUrl().trim())
                .maxCompletionTokens(ModelConstant.CUSTOM_MAX_COMPLETION_TOKENS)
                .build();
    }

    /**
     * 创建流式对话模型
     */
    public StreamingChatModel getCustomStreamingChatModel() {
        ModelConfig config = modelConfigMapper.selectOne(
                Wrappers.lambdaQuery(ModelConfig.class)
                        .eq(ModelConfig::getType, ModelType.LLM.getValue())
                        .eq(ModelConfig::getStatus, ModelConstant.IS_ACTIVE)
        );
        return new OpenAiStreamingChatModel.OpenAiStreamingChatModelBuilder()
                .strictJsonSchema(true)
                .logRequests(true)
                .logResponses(true)
                .returnThinking(config.getReturnThinking())
                .timeout(Duration.ofMinutes(3))
                .temperature(config.getTemperature())
                .apiKey(config.getApiKey().trim())
                .modelName(config.getName().trim())
                .baseUrl(config.getBaseUrl().trim())
                .maxCompletionTokens(ModelConstant.CUSTOM_MAX_COMPLETION_TOKENS)
                .build();
    }

    /**
     * 创建向量模型
     */
    public EmbeddingModel getCustomEmbeddingModel() {
        ModelConfig config = modelConfigMapper.selectOne(
                Wrappers.lambdaQuery(ModelConfig.class)
                        .eq(ModelConfig::getType, ModelType.TEXT_EMBEDDING.getValue())
                        .eq(ModelConfig::getStatus, ModelConstant.IS_ACTIVE)
        );
        return new OpenAiEmbeddingModel.OpenAiEmbeddingModelBuilder()
                .dimensions(ModelConstant.CUSTOM_DIMENSIONS)
                .maxRetries(ModelConstant.CUSTOM_MAX_RETRIES)
                .apiKey(config.getApiKey().trim())
                .modelName(config.getName().trim())
                .baseUrl(config.getBaseUrl().trim())
                .maxSegmentsPerBatch(ModelConstant.CUSTOM_MAX_SEGMENTS_PER_BATCH)
                .build();
    }

    /**
     * 创建OCR模型
     */
    public ChatModel getCustomOcrModel() {
        ModelConfig config = modelConfigMapper.selectOne(
                Wrappers.lambdaQuery(ModelConfig.class)
                        .eq(ModelConfig::getType, ModelType.OCR.getValue())
                        .eq(ModelConfig::getStatus, ModelConstant.IS_ACTIVE)
        );
        return new OpenAiChatModel.OpenAiChatModelBuilder()
                .apiKey(config.getApiKey().trim())
                .modelName(config.getName().trim())
                .baseUrl(config.getBaseUrl().trim())
                .build();
    }
}
